//----------------------------------------------------------------------------//
//                                                                            //
// ozz-animation is hosted at http://github.com/guillaumeblanc/ozz-animation  //
// and distributed under the MIT License (MIT).                               //
//                                                                            //
// Copyright (c) Guillaume Blanc                                              //
//                                                                            //
// Permission is hereby granted, free of charge, to any person obtaining a    //
// copy of this software and associated documentation files (the "Software"), //
// to deal in the Software without restriction, including without limitation  //
// the rights to use, copy, modify, merge, publish, distribute, sublicense,   //
// and/or sell copies of the Software, and to permit persons to whom the      //
// Software is furnished to do so, subject to the following conditions:       //
//                                                                            //
// The above copyright notice and this permission notice shall be included in //
// all copies or substantial portions of the Software.                        //
//                                                                            //
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR //
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   //
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    //
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER //
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING    //
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER        //
// DEALINGS IN THE SOFTWARE.                                                  //
//                                                                            //
//----------------------------------------------------------------------------//

#include "SkrAnim/ozz/ik_aim_job.h"
#include "SkrAnim/ozz/base/maths/math_ex.h"
#include "SkrAnim/ozz/base/maths/simd_quaternion.h"

#include <cassert>

using namespace ozz::math;

namespace ozz
{
namespace animation
{

bool IKAimJob::Validate() const
{
    bool valid = true;
    valid &= joint != nullptr;
    valid &= joint_correction != nullptr;
    valid &= ozz::math::AreAllTrue1(ozz::math::IsNormalizedEst3(forward));
    return valid;
}

namespace
{

// When there's an offset, the forward vector needs to be recomputed.
// The idea is to find the vector that will allow the point at offset position
// to aim at target position. This vector starts at joint position. It ends on a
// line perpendicular to pivot-offset line, at the intersection with the sphere
// defined by target position (centered on joint position). See geogebra
// diagram: media/doc/src/ik_aim_offset.ggb
bool ComputeOffsettedForward(_SimdFloat4 _forward, _SimdFloat4 _offset, _SimdFloat4 _target, SimdFloat4* _offsetted_forward)
{
    // AO is projected offset vector onto the normalized forward vector.
    assert(ozz::math::AreAllTrue1(ozz::math::IsNormalizedEst3(_forward)));
    const SimdFloat4 AOl = Dot3(_forward, _offset);

    // Compute square length of ac using Pythagorean theorem.
    const SimdFloat4 ACl2 = Length3Sqr(_offset) - AOl * AOl;

    // Square length of target vector, aka circle radius.
    const SimdFloat4 r2 = Length3Sqr(_target);

    // If offset is outside of the sphere defined by target length, the target
    // isn't reachable.
    if (AreAllTrue1(CmpGt(ACl2, r2)))
    {
        return false;
    }

    // AIl is the length of the vector from offset to sphere intersection.
    const SimdFloat4 AIl = SqrtX(r2 - ACl2);

    // The distance from offset position to the intersection with the sphere is
    // (AIl - AOl) Intersection point on the sphere can thus be computed.
    *_offsetted_forward = _offset + _forward * SplatX(AIl - AOl);

    return true;
}
} // namespace

bool IKAimJob::Run() const
{
    if (!Validate())
    {
        return false;
    }

    using math::Float4x4;
    using math::SimdFloat4;
    using math::SimdQuaternion;

    // If matrices aren't invertible, they'll be all 0 (ozz::math
    // implementation), which will result in identity correction quaternions.
    SimdInt4       invertible;
    const Float4x4 inv_joint = Invert(*joint, &invertible);

    // Computes joint to target vector, in joint local-space (_js).
    const SimdFloat4 joint_to_target_js      = TransformPoint(inv_joint, target);
    const SimdFloat4 joint_to_target_js_len2 = Length3Sqr(joint_to_target_js);

    // Recomputes forward vector to account for offset.
    // If the offset is further than target, it won't be reachable.
    SimdFloat4 offsetted_forward;
    bool       lreached = ComputeOffsettedForward(forward, offset, joint_to_target_js, &offsetted_forward);
    // Copies reachability result.
    // If offsetted forward vector doesn't exists, target position cannot be
    // aimed.
    if (reached != nullptr)
    {
        *reached = lreached;
    }

    if (!lreached ||
        AreAllTrue1(CmpEq(joint_to_target_js_len2, simd_float4::zero())))
    {
        // Target can't be reached or is too close to joint position to find a
        // direction.
        *joint_correction = SimdQuaternion::identity();
        return true;
    }

    // Calculates joint_to_target_rot_ss quaternion which solves for
    // offsetted_forward vector rotating onto the target.
    const SimdQuaternion joint_to_target_rot_js =
        SimdQuaternion::FromVectors(offsetted_forward, joint_to_target_js);

    // Calculates rotate_plane_js quaternion which aligns joint up to the pole
    // vector.
    const SimdFloat4 corrected_up_js =
        TransformVector(joint_to_target_rot_js, up);

    // Compute (and normalize) reference and pole planes normals.
    const SimdFloat4 pole_vector_js = TransformVector(inv_joint, pole_vector);
    const SimdFloat4 ref_joint_normal_js =
        Cross3(pole_vector_js, joint_to_target_js);
    const SimdFloat4 joint_normal_js =
        Cross3(corrected_up_js, joint_to_target_js);
    const SimdFloat4 ref_joint_normal_js_len2 = Length3Sqr(ref_joint_normal_js);
    const SimdFloat4 joint_normal_js_len2     = Length3Sqr(joint_normal_js);

    const SimdFloat4 denoms =
        SetZ(SetY(joint_to_target_js_len2, joint_normal_js_len2), ref_joint_normal_js_len2);

    SimdFloat4     rotate_plane_axis_js;
    SimdQuaternion rotate_plane_js;
    // Computing rotation axis and plane requires valid normals.
    if (AreAllTrue3(CmpNe(denoms, simd_float4::zero())))
    {
        const SimdFloat4 rsqrts =
            RSqrtEstNR(SetZ(SetY(joint_to_target_js_len2, joint_normal_js_len2), ref_joint_normal_js_len2));

        // Computes rotation axis, which is either joint_to_target_js or
        // -joint_to_target_js depending on rotation direction.
        rotate_plane_axis_js = joint_to_target_js * SplatX(rsqrts);

        // Computes angle cosine between the 2 normalized plane normals.
        const SimdFloat4 rotate_plane_cos_angle = Dot3(
            joint_normal_js * SplatY(rsqrts), ref_joint_normal_js * SplatZ(rsqrts)
        );
        const SimdFloat4 axis_flip =
            And(SplatX(Dot3(ref_joint_normal_js, corrected_up_js)),
                simd_int4::mask_sign());
        const SimdFloat4 rotate_plane_axis_flipped_js =
            Xor(rotate_plane_axis_js, axis_flip);

        // Builds quaternion along rotation axis.
        const SimdFloat4 one = simd_float4::one();
        rotate_plane_js      = SimdQuaternion::FromAxisCosAngle(
            rotate_plane_axis_flipped_js, Clamp(-one, rotate_plane_cos_angle, one)
        );
    }
    else
    {
        rotate_plane_axis_js = joint_to_target_js * SplatX(RSqrtEstXNR(denoms));
        rotate_plane_js      = SimdQuaternion::identity();
    }

    // Twists rotation plane.
    SimdQuaternion twisted;
    if (twist_angle != 0.f)
    {
        // If a twist angle is provided, rotation angle is rotated around joint to
        // target vector.
        const SimdQuaternion twist_ss = SimdQuaternion::FromAxisAngle(
            rotate_plane_axis_js, simd_float4::Load1(twist_angle)
        );
        twisted = twist_ss * rotate_plane_js * joint_to_target_rot_js;
    }
    else
    {
        twisted = rotate_plane_js * joint_to_target_rot_js;
    }

    // Weights output quaternion.

    // Fix up quaternions so w is always positive, which is required for NLerp
    // (with identity quaternion) to lerp the shortest path.
    const SimdFloat4 twisted_fu =
        Xor(twisted.xyzw, And(simd_int4::mask_sign(),
                              CmpLt(SplatW(twisted.xyzw), simd_float4::zero())));

    if (weight < 1.f)
    {
        // NLerp start and mid joint rotations.
        const SimdFloat4 identity    = simd_float4::w_axis();
        const SimdFloat4 simd_weight = Max0(simd_float4::Load1(weight));
        joint_correction->xyzw =
            NormalizeEst4(Lerp(identity, twisted.xyzw, simd_weight));
    }
    else
    {
        // Quaternion doesn't need interpolation
        joint_correction->xyzw = twisted_fu;
    }

    return true;
}
} // namespace animation
} // namespace ozz
